package com.zz.common.kafka

import scala.util.control.NonFatal

import akka.actor.ActorSystem
import akka.kafka.{ProducerMessage, ProducerSettings}
import akka.kafka.scaladsl.Producer
import akka.stream.scaladsl._
import akka.stream._

import com.zz.common.metrics.SimpleMetrics
import org.apache.kafka.clients.producer.ProducerRecord
import org.apache.kafka.common.serialization.Serializer

/**
  * K: Kafka key type
  * V: Kafka value type
  * T: received message's type
  */
trait ProducerPipeLine[K, V, T] {

  val producerKeySerializer: Serializer[K]
  val producerValueSerializer: Serializer[V]

  val bufferSize: Int = 1000000
  val overflowStrategy: OverflowStrategy = akka.stream.OverflowStrategy.backpressure

  def keyTransfer(t: T): K

  def valueTransfer(t: T): V

  def queue(topic: String)(implicit system: ActorSystem, mt: Materializer): SourceQueueWithComplete[T] = {
    val metrics = SimpleMetrics(system)
    Source
      .queue[T](bufferSize, overflowStrategy)
      .map(e => ProducerMessage.Message(new ProducerRecord(topic, keyTransfer(e), valueTransfer(e)), e))
      .via(Producer.flexiFlow(ProducerSettings(system, producerKeySerializer, producerValueSerializer): ProducerSettings[K, V]))
      .withAttributes(ActorAttributes.supervisionStrategy(decider))
      .map(_ => metrics.kafkaSentCounter.inc())
      .to(Sink.ignore)
      .run()
  }

  protected val decider: Supervision.Decider = {
    case NonFatal(e) =>
      e.printStackTrace()
      Supervision.Resume
  }
}
